//
// Created by twetec on 25-6-22.
//

#include "handle_icmp_probe_task.h"

#include <sys/socket.h>
#include <arpa/inet.h>
#include <netinet/ip_icmp.h>
#include <netinet/ip6.h>
#include <netinet/icmp6.h>
#include <coroutine>

#include <bs_thread_pool/BS_thread_pool.hpp>

using namespace std::chrono_literals;

namespace network_prober::probe
{
// 线程池
static BS::thread_pool s_iopool{3000};

// identifier的随机数生成器
static std::atomic_uint16_t s_identifierCounter{0};

// ICMPProbeIPv4::ICMPProbeIPv4(const std::string &interface, const std::string &target)
//     : m_interface{interface}, m_target{target}, m_iopool{1}
// {
//     m_iopool.start();
//
//     m_pingPtr = std::make_unique<asio2::ping_ipv4>(m_iopool.get(0));
//
//     // 一次ping的超时时间
//     m_pingPtr->set_timeout(13s);
//     // 两次ping之间的间隔时间
//     m_pingPtr->set_interval(20s);
//
//     // 在socket初始化时调用
//     m_pingPtr->bind_init([this]() {
//         // 绑定请求从哪个网口发送
//         asio::ip::icmp::endpoint endpoint{asio::ip::address_v4::from_string(m_interface), 0};
//         m_pingPtr->socket().bind(endpoint);
//     });
// }
//
// void ICMPProbeIPv4::probe()
// {
//     // 发起请求
//     m_pingPtr->start(m_target);
// }
//
// void ICMPProbeIPv4::setOnProbedCallback(std::function<void(std::int64_t)>&& callback)
// {
//     m_onProbedCallback = callback;
//
//     std::cout << "bind_recv: 外面" << std::this_thread::get_id() << std::endl;
//     // 请求响应时调用
//     m_pingPtr->bind_recv([this](asio2::icmpv4_rep& rep) {
//         std::cout << "bind_recv: 里面" << std::this_thread::get_id() << std::endl;
//         m_onProbedCallback((rep.is_timeout() || rep.lag.count()==-1)? -1 : std::chrono::duration_cast<std::chrono::nanoseconds>(rep.lag).count());
//         m_iopool.stop();
//     });
// }
//
//
// ICMPProbeIPv4::~ICMPProbeIPv4()
// {
//     SPDLOG_INFO("yyyy 析构了..");
//
//     std::cout << "~ICMPProbeIPv4: " << std::this_thread::get_id() << std::endl;
//
//     // m_iopool.stop();
//
//     // 停止探测
//     // m_ping.wait_stop();
//     // m_ping.destroy()
//     // m_ping.stop();
//     // m_ping.destroy();
// }

/**
 * icmp ipv4 协议头
 */
// typedef struct
// {
//     std::uint8_t type;  // 类型(8=请求，0=应答)
//     std::uint8_t code;      // 代码(通常为0)
//     std::uint16_t checksum; // 校验和
//     std::uint16_t id;       // 标识符
//     std::uint16_t seq;      // 序列号
// } ICMPHeaderIPv4;

typedef struct
{
    struct icmphdr header;  // icmp v4 协议头
    struct timeval timestamp;  // 时间戳
    char payload[56];   //  负载
} ICMPPacketIPv4;

unsigned short computeICMPChecksumIPv4(unsigned short* addr, int len)
{
    int nleft = len;    // 数据包字节长度，用于计数未经过计算的字节数
    int sum = 0;
    unsigned short* w = addr;   // 数据包指针
    unsigned short answer = 0;  // 16位

    // 计算报文中两个字节的和
    while (nleft > 1) {
        sum += *w++;    // 1、从数据包中取16位；2、将取出的16位累加到和中 3、指针再移动16位，指向下两个要累加的数据字节内存
        nleft -= 2; // 将未计算的字节数减少2个，直到只剩1或0个字节
    }

    // 如果报文长度是奇数，最后一个字节作为高8位，再用0填充1个字节(低8位)
    if (nleft == 1) {
        *(unsigned char *)(&answer) = *(unsigned char *)w;  // 把16位类型指针转为8位类型指针，取剩下的1个字节复制给answer
        sum += answer;  // 将低位累加
    }

    sum = (sum >> 16) + (sum & 0xffff); // 高16位 + 低16位 = sum只剩16位
    sum += (sum >> 16);
    answer = ~sum;
    return answer;
}

HandleICMPProbeIPv4Awaiter::HandleICMPProbeIPv4Awaiter(const database::ProbeTaskModel &probeTask)
    : m_interface{probeTask.interface}, m_target{probeTask.address}
{
}

bool HandleICMPProbeIPv4Awaiter::await_ready()
{
    return false;
}

void HandleICMPProbeIPv4Awaiter::await_suspend(std::coroutine_handle<> handle)
{
    s_iopool.submit_task([this, handle] {
        // 预设延迟未-1
        m_delay = -1;

        // 当socket套接字无效时，尝试3次
        for (std::size_t i = 0; i < 3; i ++) {
            // 创建UDP套接字
            int sockfd = socket(AF_INET, SOCK_RAW, IPPROTO_ICMP);
            if (sockfd < 0) {
                // 输出日志
                SPDLOG_INFO("sockfd创建失败: {}", strerror(errno));
                // 继续尝试获取套接字
                continue;
            }

            // 设置接收超时(3s)
            struct timeval tv;
            tv.tv_sec = 6000 / 1000;
            tv.tv_usec = (6000 % 1000) * 1000;
            if (setsockopt(sockfd, SOL_SOCKET, SO_RCVTIMEO, &tv, sizeof(tv)) < 0) {
               // 关闭socket
               close(sockfd);
               // 输出日志
               SPDLOG_INFO("sockfd: {} - 超时设置时间失败: {}", sockfd, strerror(errno));
               // 判断错误是否是套接字
               if (errno == EBADF) continue;
               // 结束循环
               break;
            }

            // 设定网口地址
            struct sockaddr_in interfaceAddr;
            memset(&interfaceAddr, 0, sizeof(interfaceAddr));
            interfaceAddr.sin_family = AF_INET;
            interfaceAddr.sin_port = htons(0);
            if (inet_pton(AF_INET, m_interface.c_str(), &interfaceAddr.sin_addr) <= 0) {
                // 关闭socket
                close(sockfd);
                // 输出日志
                SPDLOG_INFO("sockfd: {} - 网口地址设置失败: {}", sockfd, strerror(errno));
                // 判断错误是否是套接字
                if (errno == EBADF) continue;
                // 结束循环
                break;
            }

            // 绑定本机网口
            if (bind(sockfd, (sockaddr *)&interfaceAddr, sizeof(interfaceAddr)) < 0) {
                // 关闭socket
                close(sockfd);
                // 输出日志
                SPDLOG_INFO("sockfd: {} - 网口地址绑定失败: {}", sockfd, strerror(errno));
                // 判断错误是否是套接字
                if (errno == EBADF) continue;
                // 结束循环
                break;
            }

            // 设置目标地址
            struct sockaddr_in targetAddr;
            memset(&targetAddr, 0, sizeof(targetAddr));
            targetAddr.sin_family = AF_INET;
            if (inet_pton(AF_INET, m_target.c_str(), &targetAddr.sin_addr) <= 0) {
                // 关闭socket
                close(sockfd);
                // 输出日志
                SPDLOG_INFO("sockfd: {} - 目标地址设置失败: {}", sockfd, strerror(errno));
                // 判断错误是否是套接字
                if (errno == EBADF) continue;
                // 结束循环
                break;
            }

            // 填充ICMP包
            ICMPPacketIPv4 sendPacket;
            memset(&sendPacket, 0, sizeof(sendPacket));
            sendPacket.header.type = ICMP_ECHO;
            sendPacket.header.code = 0;
            sendPacket.header.un.echo.id = getpid() & 0xFFFF;
            strncpy(sendPacket.payload, "PING", sizeof(sendPacket.payload)-1);
            // 准备校验和
            sendPacket.header.un.echo.sequence = 0;
            gettimeofday(&sendPacket.timestamp, NULL);
            sendPacket.header.checksum = 0;
            sendPacket.header.checksum = computeICMPChecksumIPv4((unsigned short*)&sendPacket, sizeof(sendPacket));

            // 请求计时 - 开始时间
            std::chrono::time_point<std::chrono::steady_clock> detectStartTime = std::chrono::steady_clock::now();

            // 发送ICMP包
            if (sendto(sockfd, &sendPacket, sizeof(sendPacket), 0, (struct sockaddr *)&targetAddr, sizeof(targetAddr)) <= 0) {
                // 关闭socket
                close(sockfd);
                // 输出日志
                SPDLOG_INFO("sockfd: {} - ICMP数据包发送失败: {}", sockfd, strerror(errno));
                // 判断错误是否是套接字
                if (errno == EBADF) continue;
                // 结束循环
                break;
            }

            // 尝试接收响应
            // ip头部 + icmp头部 + icmp数据
            char recvBuf[1024];
            socklen_t targetAddrLen = sizeof(targetAddr);

            // 接收响应
            ssize_t recvdBytesCount = recvfrom(sockfd, recvBuf, sizeof(recvBuf), 0, (struct sockaddr *)&targetAddr, &targetAddrLen);

            // 关闭socket
            close(sockfd);

            // 如果未接收到数据
            if (recvdBytesCount <= 0) {
                // 输出日志
                SPDLOG_INFO("sockfd: {} - 数据包接收失败: {}", sockfd, strerror(errno));
                // 判断错误是否是套接字
                if (errno == EBADF) continue;
                // 结束执行
                break;
            }

            // 如果接收到数据，则作处理
            // 解析IP头获取ICMP包位置
            struct iphdr *ipHeader = (struct iphdr *)recvBuf;
            size_t ipHeaderLen = ipHeader->ihl * 4;
            // 如果包长度不够，则说明接收失败
            if (recvdBytesCount < ipHeaderLen + ICMP_MINLEN) {
                // 输出日志
                SPDLOG_INFO("sockfd: {} - ICMP数据包长度不够: {}", sockfd, strerror(errno));
                // 结束循环
                break;
            }

            // 提取icmp协议头 (跳过ip头)
            struct icmphdr *icmpHdr = (struct icmphdr *)(recvBuf + ipHeaderLen);
            // 检查是否是我们的响应
            if (icmpHdr->type == ICMP_ECHOREPLY && icmpHdr->un.echo.id == sendPacket.header.un.echo.id) {
                // 计算延迟
                m_delay = std::chrono::duration_cast<std::chrono::nanoseconds>(std::chrono::steady_clock::now() - detectStartTime).count();
                // 输出日志
                SPDLOG_INFO("观测任务执行成功: {} 延迟", m_delay);
                // 结束循环
                break;
            }
        }

        // 继续执行
        handle.resume();
    });
}

std::int64_t HandleICMPProbeIPv4Awaiter::await_resume()
{
    return m_delay;
}

HandleICMPProbeIPv4Awaiter::~HandleICMPProbeIPv4Awaiter()
{

}


// ICMPProbeIPv6::ICMPProbeIPv6(const std::string &interface, const std::string &target)
//     : m_interface(interface), m_target(target)
// {
//     // 一次ping的超时时间
//     m_ping.set_timeout(13s);
//     // 两次ping之间的间隔时间
//     m_ping.set_interval(20s);
//
//     // 在socket初始化时调用
//     m_ping.bind_init([this]() {
//         // 绑定请求从哪个网口发送
//         asio::ip::icmp::endpoint endpoint(asio::ip::address_v6::from_string(m_interface), 0);
//         m_ping.socket().bind(endpoint);
//     });
// }
//
// void ICMPProbeIPv6::probe()
// {
//     // 发起请求
//     m_ping.start(m_target);
// }
//
// void ICMPProbeIPv6::setOnProbedCallback(std::function<void(std::int64_t)> callback)
// {
//     m_onProbedCallback = std::move(callback);
//
//     // 返回一个指向当前对象的shared_ptr，并且计数器加了1
//     auto self{shared_from_this()};
//     // 响应时调用
//     m_ping.bind_recv([this, self](asio2::icmpv6_rep& rep) {
//         // 回调
//         m_onProbedCallback((rep.is_timeout() || rep.lag.count()==-1)? -1 : std::chrono::duration_cast<std::chrono::nanoseconds>(rep.lag).count());
//         // 停止探测
//         m_ping.stop();
//         // 该代码块执行完， self被释放，当前对象计数器减1.释放当前对象
//     });
// }
//
// ICMPProbeIPv6::~ICMPProbeIPv6()
// {
//     m_ping.destroy();
// }

HandleICMPProbeIPv6Awaiter::HandleICMPProbeIPv6Awaiter(const database::ProbeTaskModel &probeTask)
    : m_interface(probeTask.interface), m_target(probeTask.address)
{
}

bool HandleICMPProbeIPv6Awaiter::await_ready()
{
    return false;
}

typedef struct
{
    struct icmp6_hdr header;  // icmp v4 协议头
    struct timeval tv;  // 时间戳
    char payload[64-sizeof(icmp6_hdr)];   //  负载
} ICMPPacketIPv6;

void HandleICMPProbeIPv6Awaiter::await_suspend(std::coroutine_handle<> handle)
{
    s_iopool.submit_task([this, handle] {
        // 预设延迟未-1
        m_delay = -1;

        // 当socket套接字无效时，尝试3次
        for (std::size_t i = 0; i < 3; i ++) {
            // 创建UDP套接字
            int sockfd = socket(AF_INET6, SOCK_RAW, IPPROTO_ICMPV6);
            if (sockfd < 0) {
                // 输出日志
                SPDLOG_INFO("sockfd创建失败: {}", strerror(errno));
                // 继续尝试获取套接字
                continue;
            }

            // 设置接收超时(3s)
            struct timeval tv;
            tv.tv_sec = 6000 / 1000;
            tv.tv_usec = (6000 % 1000) * 1000;
            if (setsockopt(sockfd, SOL_SOCKET, SO_RCVTIMEO, &tv, sizeof(tv)) < 0) {
               // 关闭socket
               close(sockfd);
               // 输出日志
               SPDLOG_INFO("sockfd: {} - 超时设置时间失败: {}", sockfd, strerror(errno));
               // 判断错误是否是套接字
               if (errno == EBADF) continue;
               // 结束循环
               break;
            }

            // // 设定网口地址
            // struct sockaddr_in6 interfaceAddr;
            // memset(&interfaceAddr, 0, sizeof(interfaceAddr));
            // interfaceAddr.sin6_family = AF_INET6;
            // interfaceAddr.sin6_port = htons(0);
            // if (inet_pton(AF_INET6, m_interface.c_str(), &interfaceAddr.sin6_addr) <= 0) {
            //     // 关闭socket
            //     close(sockfd);
            //     // 输出日志
            //     SPDLOG_INFO("sockfd: {} - 网口地址设置失败: {}", sockfd, strerror(errno));
            //     // 判断错误是否是套接字
            //     if (errno == EBADF) continue;
            //     // 结束循环
            //     break;
            // }
            //
            // // 绑定本机网口
            // if (bind(sockfd, (sockaddr *)&interfaceAddr, sizeof(interfaceAddr)) < 0) {
            //     // 关闭socket
            //     close(sockfd);
            //     // 输出日志
            //     SPDLOG_INFO("sockfd: {} - 网口地址绑定失败: {}", sockfd, strerror(errno));
            //     // 判断错误是否是套接字
            //     if (errno == EBADF) continue;
            //     // 结束循环
            //     break;
            // }

            // 设置目标地址
            struct sockaddr_in6 targetAddr;
            memset(&targetAddr, 0, sizeof(targetAddr));
            targetAddr.sin6_family = AF_INET6;
            if (inet_pton(AF_INET6, m_target.c_str(), &targetAddr.sin6_addr) <= 0) {
                // 关闭socket
                close(sockfd);
                // 输出日志
                SPDLOG_INFO("sockfd: {} - 目标地址设置失败: {}", sockfd, strerror(errno));
                // 判断错误是否是套接字
                if (errno == EBADF) continue;
                // 结束循环
                break;
            }

            /// 填充ICMP包
            ICMPPacketIPv6 sendPacket;
            memset(&sendPacket, 0, sizeof(sendPacket));
            sendPacket.header.icmp6_type = ICMP6_ECHO_REQUEST;
            sendPacket.header.icmp6_code = 0;
            sendPacket.header.icmp6_cksum = 0;  // checksum will be calculated later
            sendPacket.header.icmp6_id = htons(getpid() & 0xFFFF);  // Identifier
            sendPacket.header.icmp6_seq = htons(1);  // Sequence number
            strncpy(sendPacket.payload, "PING", sizeof(sendPacket.payload)-1);

            // 请求计时 - 开始时间
            std::chrono::time_point<std::chrono::steady_clock> detectStartTime = std::chrono::steady_clock::now();

            // 发送ICMP包
            if (sendto(sockfd, &sendPacket, sizeof(sendPacket), 0, (struct sockaddr *)&targetAddr, sizeof(targetAddr)) <= 0) {
                // 关闭socket
                close(sockfd);
                // 输出日志
                SPDLOG_INFO("sockfd: {} - ICMP数据包发送失败: {}", sockfd, strerror(errno));
                // 判断错误是否是套接字
                if (errno == EBADF) continue;
                // 结束循环
                break;
            }

            // 尝试接收响应
            // ip头部 + icmp头部 + icmp数据
            char recvBuf[1024];
            socklen_t targetAddrLen = sizeof(targetAddr);

            // 接收响应
            ssize_t recvdBytesCount = recvfrom(sockfd, recvBuf, sizeof(recvBuf), 0, (struct sockaddr *)&targetAddr, &targetAddrLen);

            // 关闭socket
            close(sockfd);

            // 如果未接收到数据
            if (recvdBytesCount <= 0) {
                // 输出日志
                SPDLOG_INFO("sockfd: {} - ICMP数据包接收失败: {}", sockfd, strerror(errno));
                // 判断错误是否是套接字
                if (errno == EBADF) continue;
                // 结束执行
                break;
            }

            // 如果接收到数据，则作处理
            // 解析IP头获取ICMP包位置
            struct icmp6_hdr *icmpHdr = (struct icmp6_hdr *)recvBuf;
            // 检查是否是我们的响应
            if (icmpHdr->icmp6_type == ICMP6_ECHO_REPLY) {
                // 计算延迟
                m_delay = std::chrono::duration_cast<std::chrono::nanoseconds>(std::chrono::steady_clock::now() - detectStartTime).count();
                // 输出日志
                SPDLOG_INFO("观测任务执行成功: {} 延迟", m_delay);
                // 结束循环
                break;
            }
        }
        // 继续执行
        handle.resume();
    });
}

std::int64_t HandleICMPProbeIPv6Awaiter::await_resume()
{
    return m_delay;
}


}